package com.learn.algorithm.tree;

/**
 * 前缀树类
 */
public class Trie {

    private TrieNode root;

    public Trie() {
        root = new TrieNode();
    }

    public void insert(String word) {
        TrieNode node = root;
        node.pass++;
        for (int i = 0; i < word.length(); i++) {
            // 找到字符在children数组中的索引位置
            int index = word.charAt(i) - 'a';
            if (node.children[index] == null) { // 不存在当前字符则创建新节点
                node.children[index] = new TrieNode();
            }
            // 节点指向下一个节点
            node = node.children[index];
            node.pass++;
        }
        node.end++;
    }

    /**
     * 前缀树种以word单词出现的次数
     */
    public int countWordsEqualTo(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            int index = word.charAt(i) - 'a';
            if (node.children[index] == null) {
                return 0;
            } else {
                node = node.children[index];
            }
        }
        return node.end;
    }

    /**
     * 前缀树种以word为前缀的单词出现的次数
     */
    public int countWordsStartsWith(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            int index = word.charAt(i) - 'a';
            if (node.children[index] == null) {
                return 0;
            } else {
                node = node.children[index];
            }
        }
        return node.pass;
    }

    /**
     * 从前缀树种删除word
     */
    public void erase(String word) {
        if (countWordsEqualTo(word) == 0) return;
        TrieNode node = root;
        node.pass--;
        for (int i = 0; i < word.length(); i++) {
            // 找到字符在children数组中的索引位置
            int index = word.charAt(i) - 'a';
            if (--node.children[index].pass == 0) {
                node.children[index] = null;
                return;
            }
            node = node.children[index];
        }
        node.end--;
    }

    public boolean search(String word) {
        return countWordsEqualTo(word) > 0;
    }

    public boolean startsWith(String prefix) {
        return countWordsStartsWith(prefix) > 0;
    }

    public static class TrieNode {

        public TrieNode() {
            pass = 0;
            end = 0;
            children = new TrieNode[26];
        }

        /**
         * 经过当前节点的次数
         */
        public int pass;

        /**
         * 以当前节点为结尾的次数
         */
        public int end;
        public TrieNode[] children;
    }

    public static void main(String[] args) {
        Trie trie = new Trie();
        trie.insert("apple");
        trie.search("apple");
        trie.search("app");
        trie.startsWith("app");
        trie.insert("app");
        trie.search("app");
    }
}
